#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author :   Liang Kang
@Contact:   gangkanli1219@163.com
@Time   :   2018/8/1 17:58
@Desc   :
"""
import argparse
from abc import abstractmethod
from pathlib import Path

from dltools.data.dataset.cropper import VOCImageCropper
from dltools.utils.basic import is_rectangle_overlap
from dltools.utils.log import get_console_logger

from utils.file_read import GoodsReader


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--root', type=str, default='', dest='root',
        help='The directory where the image data and the annotation is stored.'
             '保存图像与标记数据的根目录, 图像与其相应标记应当在同一级目录。')
    parser.add_argument(
        '--output', type=str, default='', dest='output',
        help='the root for all the output.'
             '输出结果的根目录')
    parser.add_argument(
        '--ratio', type=float, default=0.5, dest='ratio',
        help='高与宽之比.')
    parser.add_argument(
        '--height', type=float, default=0.9, dest='height',
        help='保留高的比例.')
    parser.add_argument(
        '--size', type=int, default=640, dest='size',
        help='子图尺寸.')
    parser.add_argument(
        '--stride', type=int, default=240, dest='stride',
        help='步长.')
    return parser.parse_args()


class ShelfCropper(VOCImageCropper):

    @abstractmethod
    def _crop_image(self):
        """
        切割子图

        Returns
        -------

        """
        x, y = self._buf_data['x'], self._buf_data['y']
        w_min, w_max = x, x + self._size[1] - 1
        h_min, h_max = y, y + self._size[0] - 1
        output_objects = []
        for ob in self._image_info['objects']:
            box = ob['box']
            new_ob = ob.copy()
            if not is_rectangle_overlap(box, [h_min, w_min, h_max, w_max]):
                continue
            else:
                xmin, ymin = max(w_min, box[1]), max(h_min, box[0])
                xmax, ymax = min(w_max, box[3]), min(h_max, box[2])
                height = (ymax - ymin) > self._threshold['height'] * (
                            box[2] - box[0])
                ratio = (ymax - ymin) / (xmax - xmin) < self._threshold['ratio']
                if height and ratio:
                    new_ob['box'] = [ymin - y, xmin - x, ymax - y, xmax - x]
                    output_objects.append(new_ob)

        if len(output_objects) > 0:
            self._is_write = self._do_write
            self._buf_data['objects'] = output_objects
            self._buf_data['sub_image'] = (
                self.image[h_min:h_max + 1, w_min:w_max + 1, :])
        else:
            self._is_write = False


def crop():
    reader = GoodsReader(ARGS.root, logger)
    output = Path(ARGS.output)
    if not output.exists():
        output.mkdir()
    path = ''
    size = (ARGS.size, ARGS.size)
    stride = (ARGS.stride, ARGS.stride)
    threshold = {'height': ARGS.height, 'ratio': ARGS.ratio}
    for idx, data in enumerate(reader):
        if idx % 100 == 0:
            path = output / '{:06}'.format(idx // 100)
            if not path.exists():
                path.mkdir()
            path = str(path)
        cropper = ShelfCropper(data['image'], data['xml'],
                               path, size, stride, threshold,
                               logger=logger)
        cropper.update()


if __name__ == '__main__':
    ARGS = parse_args()
    logger = get_console_logger('Cropper')
    crop()
